#!/usr/bin/env python3
"""
Script: modify_topas_histories.py
Function: Batch modify NumberOfHistoriesInRun parameter in carbon.txt files
          across dedx_m1 to dedx_m25 directories
Author: Converted from bash script
"""

import os
import re
import shutil
import sys
from datetime import datetime
from pathlib import Path


def main():
    # Colors for output (Windows compatible)
    RED = '\033[0;31m'
    GREEN = '\033[0;32m'
    YELLOW = '\033[1;33m'
    BLUE = '\033[0;34m'
    NC = '\033[0m'  # No Color

    success_count = 0
    error_count = 0
    total_files = 0

    print(f"{BLUE}==================================================={NC}")
    print(f"{BLUE}  TOPAS Parameter Batch Modification Script{NC}")
    print(f"{BLUE}==================================================={NC}")
    print(f"{YELLOW}Task: Modify NumberOfHistoriesInRun from 500000 to 5000000{NC}")
    print(f"{YELLOW}Target directories: dedx_m1 to dedx_m25{NC}")
    print(f"{YELLOW}Target file: carbon.txt{NC}")
    print()

    # Check if dedx_m1 directory exists
    if not os.path.exists("dedx_m1"):
        print(f"{RED}Error: dedx_m1 directory not found in current directory{NC}")
        print(f"{YELLOW}Please run this script in the directory containing dedx_m1 to dedx_m25{NC}")
        sys.exit(1)

    # Create backup directory
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_dir = f"backup_{timestamp}"
    os.makedirs(backup_dir, exist_ok=True)
    print(f"{BLUE}Creating backup directory: {backup_dir}{NC}")

    # Process directories dedx_m1 to dedx_m25
    for i in range(1, 26):
        dir_name = f"dedx_m{i}"
        carbon_file = os.path.join(dir_name, "carbon.txt")

        print(f"\n{BLUE}Processing directory: {dir_name}{NC}")

        # Check if directory exists
        if not os.path.exists(dir_name):
            print(f"{YELLOW}  Warning: Directory {dir_name} does not exist, skipping{NC}")
            continue

        # Check if carbon.txt file exists
        if not os.path.exists(carbon_file):
            print(f"{YELLOW}  Warning: File {carbon_file} does not exist, skipping{NC}")
            continue

        total_files += 1

        # Create backup
        backup_file = os.path.join(backup_dir, f"{dir_name}_carbon.txt")
        shutil.copy2(carbon_file, backup_file)
        print(f"{GREEN}  ✓ Backed up to: {backup_file}{NC}")

        # Read the file content
        try:
            with open(carbon_file, 'r', encoding='utf-8') as f:
                content = f.read()
        except Exception as e:
            print(f"{RED}  ✗ Failed to read {carbon_file}: {e}{NC}")
            error_count += 1
            continue

        # Check if the target parameter exists and replace it
        pattern = r'i:So/Demo/NumberOfHistoriesInRun\s*=\s*500000'
        if re.search(pattern, content):
            # Replace 500000 with 5000000
            new_content = re.sub(pattern, 'i:So/Demo/NumberOfHistoriesInRun = 5000000', content)

            try:
                with open(carbon_file, 'w', encoding='utf-8') as f:
                    f.write(new_content)
                print(f"{GREEN}  ✓ Successfully modified: {carbon_file}{NC}")
                print(f"{GREEN}    500000 → 5000000{NC}")
                success_count += 1
            except Exception as e:
                print(f"{RED}  ✗ Failed to write {carbon_file}: {e}{NC}")
                error_count += 1
        else:
            print(f"{YELLOW}  ⚠ Target parameter not found, checking file content{NC}")
            # Show possible matching lines
            lines = content.split('\n')
            matching_lines = [f"  Line {j+1}: {line.strip()}" for j, line in enumerate(lines)
                            if 'NumberOfHistoriesInRun' in line]
            for line in matching_lines[:3]:
                print(line)
            error_count += 1

    # Display summary
    print(f"\n{BLUE}==================================================={NC}")
    print(f"{BLUE}  Modification Summary{NC}")
    print(f"{BLUE}==================================================={NC}")
    print(f"{GREEN}Successfully modified: {success_count} files{NC}")
    print(f"{RED}Failed/skipped: {error_count} files{NC}")
    print(f"{BLUE}Total processed: {total_files} files{NC}")
    print(f"{BLUE}Backup location: {backup_dir}{NC}")

    # Provide verification command
    print(f"\n{YELLOW}Verify modifications with:{NC}")
    print(f"{YELLOW}grep -r \"NumberOfHistoriesInRun\" dedx_m*/carbon.txt{NC}")

    # Provide restore command if there were errors
    if error_count > 0:
        print(f"\n{YELLOW}To restore from backup:{NC}")
        print(f"{YELLOW}for i in range(1, 26):")
        print(f"    backup_file = f'{backup_dir}/dedx_m{{i}}_carbon.txt'")
        print("    if os.path.exists(backup_file):")
        print(f"        shutil.copy2(backup_file, f'dedx_m{{i}}/carbon.txt'){NC}")

    print(f"\n{GREEN}Script execution completed!{NC}")


if __name__ == "__main__":
    main()
